# 首先导入所需第三方库
from langchain.document_loaders import UnstructuredFileLoader
from langchain.document_loaders import UnstructuredMarkdownLoader
from langchain.document_loaders import UnstructuredPDFLoader
# from langchain.document_loaders import PyPDFLoader
from langchain.document_loaders import PythonLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from tqdm import tqdm
import os

from config.index import (
    file_dir,
    Persist_dir,
    EMBEDDING_DEVICE,
    EMBEDDING_MODEL_NAME
)
# 获取文件路径函数
def get_files(dir_path):
    # args：dir_path，目标文件夹路径
    file_list = []
    for filepath, dirnames, filenames in os.walk(dir_path):
        # os.walk 函数将递归遍历指定文件夹
        for filename in filenames:
            # 通过后缀名判断文件类型是否满足要求
            if filename.endswith(".md"):
                # 如果满足要求，将其绝对路径加入到结果列表
                file_list.append(os.path.join(filepath, filename))
            elif filename.endswith(".txt"):
                file_list.append(os.path.join(filepath, filename))
            elif filename.endswith(".pdf"):
                file_list.append(os.path.join(filepath, filename))
            elif filename.endswith(".prisma"):
                file_list.append(os.path.join(filepath, filename))
            elif filename.endswith(".yaml"):
                file_list.append(os.path.join(filepath, filename))
    return file_list

# 加载文件函数
def get_text(dir_path):
    # args：dir_path，目标文件夹路径
    # 首先调用上文定义的函数得到目标文件路径列表
    file_lst = get_files(dir_path)
    # docs 存放加载之后的纯文本对象
    docs = []
    # 遍历所有目标文件
    for one_file in tqdm(file_lst):
        file_type = one_file.split('.')[-1]
        if file_type == 'md':
            loader = UnstructuredMarkdownLoader(one_file)
            docs.extend(loader.load())
        elif file_type == 'txt':
            loader = UnstructuredFileLoader(one_file)
            docs.extend(loader.load())
        elif file_type == 'pdf':
            #  mode="paged",
            # strategy='fast',
            # mode="elements",
            loader = UnstructuredPDFLoader(one_file, languages = ['chi_sim', 'eng'])
            docs.extend(loader.load()) 
        elif file_type == 'prisma':
            #  mode="paged",
            # strategy='fast',
            # mode="elements",
            loader = PythonLoader(one_file)
            print('prisma docs=========================', loader.load())
            docs.extend(loader.load()) 
        elif file_type == 'yaml':
            #  mode="paged",
            # strategy='fast',
            # mode="elements",
            loader = PythonLoader(one_file)
            print('yaml docs=========================', loader.load())
            docs.extend(loader.load()) 
        else:
            # 如果是不符合条件的文件，直接跳过
            continue
    return docs



# 加载文本文件
docs = []
for dir_path in file_dir:
    docs.extend(get_text(dir_path))
# 对文本进行分块
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200)
split_docs = text_splitter.split_documents(docs)

# 加载开源词向量模型
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": EMBEDDING_DEVICE})

# 构建向量数据库
# 加载数据库
vectordb = Chroma.from_documents(
    documents=split_docs,
    embedding=embeddings,
    persist_directory=Persist_dir  # 允许我们将persist_directory目录保存到磁盘上
)
# 将加载的向量数据库持久化到磁盘上
vectordb.persist()